import numpy as np
import sys
import dpdata, deepmd

object_dir="/home/nxu/work_dirs/test/test_dpgen/test_data_set/dp_test/deepmd/Si64"
force_file="params.pb"

dp=deepmd.infer.DeepPot(force_file)
system = dpdata.LabeledSystem(object_dir,fmt='deepmd/npy')
labeled_sys= system.predict(dp)
np.savetxt("E.txt",np.array([labeled_sys['energies'].flatten(),system['energies']]).T)
np.savetxt("f.txt",np.concatenate((labeled_sys['forces'].reshape(-1,3),system['forces'].reshape(-1,3)),axis=1))
np.savetxt("virial.txt",np.concatenate((labeled_sys['virials'].reshape(-1,9),system['virials'].reshape(-1,9)),axis=1))
data=np.loadtxt("E.txt")
rmse_E = np.sqrt(np.mean((data[:,0]-data[:,1])**2))  
data=np.loadtxt("f.txt")
data0=np.sqrt(np.sum(data[:,:3]**2,axis=1))
data1=np.sqrt(np.sum(data[:,3:]**2,axis=1))
rmse_f = np.sqrt(np.mean((data0-data1)**2))
data=np.loadtxt("virial.txt")
rmse_virial = np.sqrt(np.mean((data[:,:9]-data[:,9:])**2))
print("Energy RMSE        : %f eV" %(rmse_E))
print("Force  RMSE        : %f eV/A" %(rmse_f))
print("Virial RMSE        : %f eV" %(rmse_virial))



